import torch
from torch import nn
from torchvision.models import mobilenet_v2


class MobileNetCIFAR(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = mobilenet_v2(pretrained=False)

        # 修改首层卷积（网页8）
        self.model.features[0][0] = nn.Conv2d(
            3, 32, kernel_size=3, stride=1, padding=1, bias=False
        )

        # 调整分类层（网页6）
        self.model.classifier[1] = nn.Linear(self.model.last_channel, 10)

        # 添加形状调试
        self.shape_debug = False

    def forward(self, x):
        # if self.shape_debug:
        #     print("输入形状:", x.shape)  # 调试输出

        x = self.model.features(x)
        x = nn.functional.adaptive_avg_pool2d(x, (1, 1))
        x = torch.flatten(x, 1)
        return self.model.classifier(x)
